# -*- coding: utf-8 -*-
# @Time    : 2024/11/18 16:04
# @Author  : sjh
# @Site    : 
# @File    : scoreAction.py
# @Comment :
import json
from utils.CalculateUtils import CalcuUtils as util
import numpy as np
import matplotlib.pyplot as plt
from utils.genaral import plot_peak, plot_peak_1d, butterworth_filter_1d, smooth_filter_1d, \
    multi_frame_check_with_indices, savgol_filter_3d
from scipy.signal import find_peaks, savgol_filter

from utils.oneEuro import Filter3D

FPS = 30


class ActionScore():
    Const_KEY_POINTS = ['nose', 'left_eye', 'right_eye', 'left_ear', 'right_ear',
                        'right_ear', 'right_shoulder', 'left_elbow', 'right_elbow', 'left_wrist',
                        'right_wrist', 'left_hip', 'right_hip', 'left_knee', 'right_knee',
                        'left_ankle', 'right_ankle', 'head', 'neck', 'hip',
                        'left_big_toe', 'right_big_toe', 'left_small_toe', 'right_small_toe', 'left_heel',
                        'right_heel']

    def __init__(self):
        self.smooth = Filter3D(num=30)

    def filter_keypoints(self, keypoints, key_remove):
        if key_remove == 'all':
            key_remove = list(range(26))
        filtered_keypoints = []
        self.ori_points_list = key_remove
        keypoints_names = [self.Const_KEY_POINTS[key] for key in key_remove]
        previous_frame = None  # 用于存储上一帧的值
        keypoints_num = len(key_remove)
        for d_idx, d in enumerate(keypoints):
            remove = False
            find_num = 0
            for key, value in enumerate(d):
                if key in key_remove:
                    value = list(value.values())[0]
                    if any(v == 0.0 for v in value):
                        if previous_frame is not None:
                            # 使用上一帧的值替换当前帧
                            d[key] = previous_frame[key]

                        find_num += 1
                        if find_num == keypoints_num:
                            break
                        # break

            filtered_keypoints.append(d)
            previous_frame = d  # 更新上一帧为当前帧

        return filtered_keypoints

    def initService(self, points_list, key_filter=False):
        self.points_list = points_list
        self.points_list = [d for d in points_list if d]
        # 使用过滤函数
        if key_filter:
            self.points_list = self.filter_keypoints(self.points_list, key_filter)
        return list(range(len(points_list)))

    def getLimbsPosition(self, keyPointInt, numFrame, ndim=3):

        if ndim == 3:
            keyPointInt = int(keyPointInt)

            keypoint_name, keypoint_coords = list(self.points_list[numFrame][keyPointInt].items())[0]
            position = np.array(keypoint_coords) * 0.001
        return position

    def get_peaks(slef, data, height=None, distance=20, width=2, prominence=0.1):
        """
        使用 find_peaks 找到峰值点，并计算对应的上升点。
        Args:
            data (list or np.ndarray): 数据序列。
            peak_threshold (float): 峰值的最低阈值。
            distance (int): 相近点的距离阈值，只保留最大的点。

        Returns:
            list: 每个峰值点之前的上升点索引。
            list: 峰值点索引。
        """
        data = np.array(data)

        # 使用 find_peaks 找到峰值点
        """
        height 说明：控制峰值的高度，支持设置最小值和最大值（如果仅提供一个值，则为最小值）。 
        threshold 说明：控制峰值与其相邻样本的最小垂直距离差。 过滤掉幅度变化不明显的“伪峰值”。
        distance 说明：控制相邻峰值的最小水平距离（以样本点数为单位）。避免过于接近的峰值，通过优先保留更高的峰值。
        prominence 说明：控制峰值的“突出度”（即与周围最低谷之间的高度差）。 过滤掉周围起伏不大的峰值点。
        width 说明：控制峰值的宽度（以采样点为单位），可以设置最小和最大值。 只需要较宽的峰值（例如周期信号中的宽峰）。
        wlen 说明：设置用于计算峰值突出度或宽度的窗口长度。 如果信号较长且特定区域的特征明显，可以限定计算范围。
        rel_height 说明：用于宽度计算的相对高度比例。 控制宽度计算时，按相对高度的比例选择两个交点。
        plateau_size 说明：控制平顶峰值的平坦区域大小，可以设置最小和最大值。 当需要分析平顶波形信号时。
        """
        if height:
            peaks, properties = find_peaks(data, distance=distance, width=width, prominence=prominence, height=height)
        else:
            peaks, properties = find_peaks(data, distance=distance, width=width, prominence=prominence)

        if peaks is []:
            return [0], [0], [0]
        rising_points = []
        # 根据找到的峰值点计算上升点
        search_min_peaks = []
        diff = np.diff(data)
        search_window = 20
        for idx, peak in enumerate(peaks):
            # 在峰值点前的窗口中找到最低点
            if idx == 0:
                # min_index = np.argmin(data[0:peaks[idx]])
                search_window = int((peaks[idx] - 0) / 2)
            else:
                # min_index = np.argmin(data[peaks[idx-1]:peaks[idx]]) + peaks[idx-1]
                search_window = int((peaks[idx] - peaks[idx - 1]) / 2)
            # search_peak = int((peak + min_index)/2)
            start_idx = max(0, peak - search_window)
            local_min_idx = start_idx + np.argmin(data[start_idx:peak])
            search_min_peaks.append(local_min_idx)

            # 从最低点向前找到上升点（梯度由负变正）
            for j in range(local_min_idx - 1, -1, -1):
                if j > 0 and diff[j - 1] < 0 and diff[j] > 0:
                    rising_points.append(j)
                    break
            if idx == len(peaks) - 1:
                # peak后最后一个
                search_min_peaks.append(peak + np.argmin(data[peak:]))
        return rising_points, peaks.tolist(), search_min_peaks

    # 6m步速测试
    def action_1_score(self, jsonPath):
        def calculate_motion_data(left_ankle_3d, right_ankle_3d, fps, neck, *points):
            """
            计算脚部关键点的速度和位移，并判断开始运动的时刻。

            Args:
                left_ankle_3d (numpy.ndarray): 左脚踝关键点，形状为 (frames, 3)。
                right_ankle_3d (numpy.ndarray): 右脚踝关键点，形状为 (frames, 3)。
                fps (int): 视频的帧率。
                speed_threshold (float): 判断开始运动的速度阈值。

            Returns:
                dict: 包含速度、位移、以及开始运动时刻的信息。
            """
            # 初始化结果
            left_speeds = []
            right_speeds = []
            left_displacements = []
            right_displacements = []

            start_frame = None
            frame_count = 0  # 计数器，用于连续帧计数
            frame_interval_threshold = 15  # 连续帧间隔的阈值
            z_threshold = 10  # z 轴差值的阈值
            # 遍历所有帧，计算速度和位移
            a = []
            neck_Forward_z_dis, neck_Backward_z_dis = [], []
            vector0 = left_ankle_3d[0] - right_ankle_3d[0]
            vector0 = util.vectorProjectionHorizontal(vector0)

            for i in range(len(left_ankle_3d)):
                # 找左右脚角度
                vector = left_ankle_3d[i] - right_ankle_3d[i]
                vector = util.vectorProjectionHorizontal(vector)
                angle = util.vectorAngle(vector0, vector)
                a.append(angle)
                # 计算左右脚与第一帧高度差
                left_displacements.append(abs(left_ankle_3d[i][1] - left_ankle_3d[1][1]))
                right_displacements.append(abs(right_ankle_3d[i][1] - right_ankle_3d[1][1]))

                # 计算膝踝脚夹角
                left_vector = util.vectorAngle(util.vectorProjectionHorizontal(
                    points[2][i] - left_ankle_3d[i]), util.vectorProjectionHorizontal(points[0][i] - left_ankle_3d[i]))
                right_vector = util.vectorAngle(util.vectorProjectionHorizontal(
                    points[3][i] - left_ankle_3d[i]), util.vectorProjectionHorizontal(points[0][i] - left_ankle_3d[i]))
                # 计算与第一帧的深度差值
                neck_Forward_z_dis.append(neck[i][2] - neck[0][2])
                neck_Backward_z_dis.append(neck[i][2] - neck[-1][2])

            a = butterworth_filter_1d(np.array(a))
            neck_Forward_z_dis = multi_frame_check_with_indices(np.array(neck_Forward_z_dis))
            neck_Backward_z_dis = multi_frame_check_with_indices(np.array(neck_Backward_z_dis))
            # 计算开始运动的时间（秒）
            start_time = start_frame / fps if start_frame is not None else None

            return {
                "left_speed": left_speeds,
                "right_speed": right_speeds,
                "left_displacement": left_displacements,
                "right_displacement": right_displacements,
                "start_frame": start_frame,
                "start_time": start_time,
                "a": a,
                "neck_Forward_z_dis": neck_Forward_z_dis,
                "neck_Backward_z_dis": neck_Backward_z_dis
            }

        numFrame = self.initService(jsonPath, key_filter='all')
        left_ankle, right_ankle = [], []
        neck = []
        right_big_toe, left_big_toe = [], []
        right_knee, left_knee = [], []
        for i in numFrame:
            left_ankle.append(self.getLimbsPosition(7, i))
            right_ankle.append(self.getLimbsPosition(0, i))
            neck.append(self.getLimbsPosition(19, i))
            left_big_toe.append(self.getLimbsPosition(20, i))
            right_big_toe.append(self.getLimbsPosition(21, i))
            left_knee.append(self.getLimbsPosition(13, i))
            right_knee.append(self.getLimbsPosition(14, i))
        plot_peak(neck)
        # plot_peak(left_big_toe)

        motion_data = calculate_motion_data(left_ankle, right_ankle, FPS, neck, *(
            left_big_toe, right_big_toe, left_knee, right_knee, neck))

        plot_peak_1d(motion_data["neck_Forward_z_dis"])
        plot_peak_1d(motion_data["neck_Backward_z_dis"])

        # 找tstart和tend时刻
        tstart = next((i for i, z in enumerate(motion_data["neck_Forward_z_dis"]) if abs(z) > 0.1), None)
        tend = [i for i, z in enumerate(motion_data["neck_Backward_z_dis"]) if abs(z) < 0.1][-1]
        print('计算的时刻:', tstart, tend)
        z_tstart, z_tend = neck[tstart][2], neck[tend][2]
        distance = abs(z_tend - z_tstart)
        time = (tend - tstart) / FPS
        speed = distance / time if time > 0 else 0

        print("步速:m/s", speed)
        print("步chang:m", distance)

        """
        步速≥1.0m/s   正常步速＜1.0m/s   步速减慢
        """
        graph_dict = {}
        if speed >= 1:
            graph_dict['score'] = '正常'
        elif speed < 1:
            graph_dict['score'] = '步速减慢'
        return graph_dict

    # 5次起坐试验测试
    def action_2_score(self, jsonPath):
        numFrame = self.initService(jsonPath, key_filter='all')
        left_ankle, right_ankle = [], []
        neck = []
        right_big_toe, left_big_toe = [], []
        right_knee, left_knee = [], []
        hip = []
        head = []
        for i in numFrame:
            left_ankle.append(self.getLimbsPosition(7, i))
            right_ankle.append(self.getLimbsPosition(0, i))
            neck.append(self.getLimbsPosition(18, i))
            left_big_toe.append(self.getLimbsPosition(20, i))
            right_big_toe.append(self.getLimbsPosition(21, i))
            left_knee.append(self.getLimbsPosition(13, i))
            right_knee.append(self.getLimbsPosition(14, i))
            hip.append(self.getLimbsPosition(19, i))
            head.append(self.getLimbsPosition(17, i))
        head_y_dis = []
        for i in range(len(hip)):
            head_y_dis.append(head[i][1] - head[0][1])
        plot_peak(head)
        head_y_dis_ = np.array(head_y_dis) * -1
        head_y_dis_ = savgol_filter(np.array(head_y_dis_), 10, 4)
        # head_y_dis = butterworth_filter_1d(head_y_dis)

        rising_points, peaks, search_min_peaks = self.get_peaks(head_y_dis_, height=0.2)

        plt.plot(head_y_dis)
        plt.scatter(peaks, [head_y_dis[i] for i in peaks], color='blue', label='Peaks')
        plt.scatter(search_min_peaks, [head_y_dis[i] for i in search_min_peaks], color='red', label='search_peaks')
        plt.legend()
        plt.show()

        sit_up_time = (search_min_peaks[-1] - search_min_peaks[0]) / FPS  # s
        sit_up_times = len(peaks)
        print(peaks)
        graph_dict = {}
        graph_dict['sit_up_time'] = sit_up_time
        graph_dict['sit_up_times'] = sit_up_times
        if sit_up_time < 12:
            graph_dict['score'] = '正常'
        else:
            graph_dict['score'] = '跌倒风险较高'

        print(graph_dict)
        return graph_dict

    # 起立-行走计时试验（TUG）
    def action_3_score(self, jsonPath):
        numFrame = self.initService(jsonPath, key_filter='all')
        left_ankle, right_ankle = [], []
        neck = []
        right_big_toe, left_big_toe = [], []
        right_knee, left_knee = [], []
        hip = []
        for i in numFrame:
            left_ankle.append(self.getLimbsPosition(7, i))
            right_ankle.append(self.getLimbsPosition(0, i))
            neck.append(self.getLimbsPosition(18, i))
            left_big_toe.append(self.getLimbsPosition(20, i))
            right_big_toe.append(self.getLimbsPosition(21, i))
            left_knee.append(self.getLimbsPosition(13, i))
            right_knee.append(self.getLimbsPosition(14, i))
            hip.append(self.getLimbsPosition(19, i))
        hip_y_dis = []
        plot_peak(neck)
        plot_peak(hip)
        for i in range(len(hip)):
            hip_y_dis.append(hip[i][1] - hip[0][1])
        # hip_y_dis = multi_frame_check_with_indices(np.array(hip_y_dis))

        plot_peak_1d(hip_y_dis)
        tstart = [i for i, z in enumerate(hip_y_dis) if abs(z) > 0.05][0]
        tend = [i for i, z in enumerate(hip_y_dis) if abs(z) > 0.05][-1]
        tstart -= 5
        tend += 5
        if tstart < 0:
            tstart = 0
        if tend >= len(numFrame):
            tend = len(numFrame) - 1
        walk_time = int((tend - tstart) / FPS)
        print("开始帧、结束帧、耗时", tstart, tend, walk_time)

        # 评分标准
        graph_dict = {}
        if walk_time < 10:
            graph_dict["score"] = '可自由活动'
        elif 10 <= walk_time < 20:
            graph_dict["score"] = "大部分可独立活动"
        elif 20 <= walk_time < 30:
            graph_dict["score"] = "活动不稳定"
        else:
            graph_dict["score"] = "存在活动障碍"
        if hip[tend][2] - hip[1][2] > 0.5:
            graph_dict["score"] = "存在活动障碍"

        print(graph_dict)
        return graph_dict

    # 平衡测试
    def action_4_score_1(self, jsonPath, type=0):
        def cal_3d_cente(keypoints):
            weights = {
                'neck': 0.1,  # 脖子权重
                'hip': 0.3,  # 髋中心权重
                'left_shoulder': 0.05,
                'right_shoulder': 0.05,
                'left_hip': 0.1,  # 左髋权重
                'right_hip': 0.1,  # 右髋权重
                'left_knee': 0.1,  # 左膝权重
                'right_knee': 0.1,  # 右膝权重
                'left_foot': 0.15,  # 左脚权重
                'right_foot': 0.15  # 右脚权重
            }
            total_weight = 0
            x_sum, y_sum, z_sum = 0, 0, 0

            for point, coords in keypoints.items():
                if point in weights:
                    weight = weights[point]
                    x_sum += weight * coords[0]  # 累加 x 坐标的加权和
                    y_sum += weight * coords[1]  # 累加 y 坐标的加权和
                    z_sum += weight * coords[2]  # 累加 z 坐标的加权和
                    total_weight += weight  # 累加权重

            # 计算重心坐标
            x_center = x_sum / total_weight
            y_center = y_sum / total_weight
            z_center = z_sum / total_weight

            return x_center, y_center, z_center

        numFrame = self.initService(jsonPath, key_filter='all')
        left_ankle, right_ankle = [], []
        neck = []
        right_big_toe, left_big_toe = [], []
        right_knee, left_knee = [], []
        hip = []
        head = []
        for i in numFrame:
            left_ankle.append(self.getLimbsPosition(7, i))
            right_ankle.append(self.getLimbsPosition(0, i))
            neck.append(self.getLimbsPosition(18, i))
            left_big_toe.append(self.getLimbsPosition(20, i))
            right_big_toe.append(self.getLimbsPosition(21, i))
            left_knee.append(self.getLimbsPosition(13, i))
            right_knee.append(self.getLimbsPosition(14, i))
            hip.append(self.getLimbsPosition(19, i))
            head.append(self.getLimbsPosition(17, i))
        center_of_mass = []
        for idx in numFrame:
            keypoints = {
                'neck': neck[idx],  # 脖子权重
                'hip': hip[idx],  # 髋中心权重
                'left_shoulder': self.getLimbsPosition(5, idx),
                'right_shoulder': self.getLimbsPosition(6, idx),
                'left_hip': self.getLimbsPosition(11, idx),  # 左髋权重
                'right_hip': self.getLimbsPosition(12, idx),  # 右髋权重
                'left_knee': left_knee[idx],  # 左膝权重
                'right_knee': right_knee[idx],  # 右膝权重
                'left_ankle': left_ankle[idx],  # 左脚权重
                'right_ankle': right_ankle[idx]  # 右脚权重
            }
            center_of_mass.append(cal_3d_cente(keypoints))
        center_of_mass = savgol_filter_3d(np.array(center_of_mass))
        plot_peak(center_of_mass)
        for idx in range(1, len(numFrame) - 1):
            Diff_data = 0
            Diff_data += abs(center_of_mass[idx][0] - center_of_mass[0][0])
            Diff_data += abs(center_of_mass[idx][1] - center_of_mass[0][1])
            Diff_data += abs(center_of_mass[idx][2] - center_of_mass[0][2])
            if Diff_data > 0.10:
                tend = idx
                break
            else:
                tend = idx
        Equilibrium_time = tend / FPS
        all_time = round(len(numFrame) / FPS, 1)
        graph_dict = {}
        graph_dict['Equilibrium_time'] = Equilibrium_time
        graph_dict['all_time'] = all_time
        if Equilibrium_time < 10:
            graph_dict['score'] = 0
        elif Equilibrium_time >= 10:
            graph_dict['score'] = 1

        # 双脚并拢站立
        if type == 0:
            pass
        elif type == 1:
            pass
        elif type == 2:
            pass
        print(graph_dict)
        return graph_dict
    # 平衡测试
    def action_4_score(self, jsonPath, type=0):

        numFrame = self.initService(jsonPath, key_filter='all')
        left_ankle, right_ankle = [], []
        for i in numFrame:
            left_ankle.append(self.getLimbsPosition(7, i))
            right_ankle.append(self.getLimbsPosition(0, i))
        Diff_value_list = []
        for i in range(len(numFrame)):

            Diff_value_list.append(np.linalg.norm(left_ankle[i] - left_ankle[0]) + np.linalg.norm(right_ankle[i] - right_ankle[0]))

        Diff_value_list = smooth_filter_1d(np.array(Diff_value_list))
        plot_peak_1d(Diff_value_list)
        tend = [i for i, v in enumerate(Diff_value_list) if abs(v) > 0.15]
        if len(tend) == 0:
            tend = len(Diff_value_list)
        else:
            tend = tend[0]
        Equilibrium_time = round(tend / FPS,1)
        all_time = round(len(numFrame) / FPS, 1)
        graph_dict = {}
        graph_dict['Equilibrium_time'] = Equilibrium_time
        graph_dict['all_time'] = all_time
        if Equilibrium_time < 10:
            graph_dict['score'] = 0
        elif Equilibrium_time >= 10:
            graph_dict['score'] = 1

        # 双脚并拢站立
        if type == 0:
            pass
        elif type == 1:
            pass
        elif type == 2:
            pass
        print(graph_dict)
        return graph_dict
    # 步行速度测试 4m
    def action_5_score(self, jsonPath):
        def calculate_motion_data(left_ankle_3d, right_ankle_3d, fps, neck, *points):
            a = []
            neck_Forward_z_dis, neck_Backward_z_dis = [], []
            for i in range(len(left_ankle_3d)):
                # 计算与第一帧的深度差值
                neck_Forward_z_dis.append(neck[i][2] - neck[0][2])
                neck_Backward_z_dis.append(neck[i][2] - neck[-1][2])
            neck_Forward_z_dis = multi_frame_check_with_indices(np.array(neck_Forward_z_dis))
            neck_Backward_z_dis = multi_frame_check_with_indices(np.array(neck_Backward_z_dis))
            return {
                "neck_Forward_z_dis": neck_Forward_z_dis,
                "neck_Backward_z_dis": neck_Backward_z_dis
            }

        numFrame = self.initService(jsonPath, key_filter='all')
        left_ankle, right_ankle = [], []
        hip = []
        right_big_toe, left_big_toe = [], []
        right_knee, left_knee = [], []
        for i in numFrame:
            left_ankle.append(self.getLimbsPosition(7, i))
            right_ankle.append(self.getLimbsPosition(0, i))
            hip.append(self.getLimbsPosition(19, i))
            left_big_toe.append(self.getLimbsPosition(20, i))
            right_big_toe.append(self.getLimbsPosition(21, i))
            left_knee.append(self.getLimbsPosition(13, i))
            right_knee.append(self.getLimbsPosition(14, i))
        # plot_peak(right_big_toe)
        # plot_peak(left_big_toe)

        motion_data = calculate_motion_data(left_ankle, right_ankle, FPS, hip, *(
            left_big_toe, right_big_toe, left_knee, right_knee, hip))

        plot_peak_1d(motion_data["neck_Forward_z_dis"])
        plot_peak_1d(motion_data["neck_Backward_z_dis"])

        # 找tstart和tend时刻
        tstart = next((i for i, z in enumerate(motion_data["neck_Forward_z_dis"]) if abs(z) > 0.1), None)
        tend = [i for i, z in enumerate(motion_data["neck_Backward_z_dis"]) if abs(z) < 0.1][-1]
        tstart -= 5
        tend += 5
        if tstart < 0:
            tstart = 0
        if tend >= len(numFrame):
            tend = len(numFrame) - 1
        print('计算的时刻:', tstart, tend)
        z_tstart, z_tend = hip[tstart][2], hip[tend][2]
        walking_distance = abs(z_tend - z_tstart)
        walking_time = (tend - tstart) / FPS
        walking_speed = walking_distance / walking_time if walking_time > 0 else 0

        print("步速:m/s", walking_speed)
        graph_dict = {}
        graph_dict['walking_speed'] = walking_speed
        graph_dict['walking_distance'] = walking_distance
        graph_dict['walking_time'] = walking_time
        if walking_distance < 4:
            graph_dict['score'] = 0
        else:
            if 8.7 <= walking_time:
                graph_dict['score'] = 1
            elif 6.21 <= walking_time < 8.7:
                graph_dict['score'] = 2
            elif 4.82 <= walking_time < 6.21:
                graph_dict['score'] = 3
            elif walking_time < 4.82:
                graph_dict['score'] = 4
        print(graph_dict)
        """
        步速≥1.0m/s   正常步速＜1.0m/s   步速减慢
        """
        return graph_dict


if __name__ == '__main__':
    asc = ActionScore()
    jsonPath = r'output/action_4_score/keypoints_20241121_145640.json'
    with open(jsonPath) as f:
        data = json.load(f)
        asc.action_4_score(data)
